from torch import nn
import torch
class net(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.net1=nn.Linear(1,2).cuda()
        self.net2=nn.Linear(2,1).cuda()
    
    def forward(self,input):
        a=self.net1(input)
        c=a.to('cuda:1')
        print('c',c)
        return self.net2(c)

# n=net()
# print(n.state_dict())
# opt = torch.optim.SGD(n.parameters(), lr=0.01)
# inputs=torch.rand((2,1)).cuda()
# print(inputs)
# out=n(inputs)
# loss = out.sum()
# loss.backward()
# opt.step()
# print(n.state_dict())


inputs=torch.rand((2,1))
m=nn.GELU()
print(inputs)
print(m(inputs))